""" Vision Transformer (ViT) in PyTorch

A PyTorch implement of Vision Transformers as described in
'An Image Is Worth 16 x 16 Words: Transformers for Image Recognition at Scale' - https://arxiv.org/abs/2010.11929

The official jax code is released and available at https://github.com/google-research/vision_transformer

Acknowledgments:
* The paper authors for releasing code and weights, thanks!
* I fixed my class token impl based on Phil Wang's https://github.com/lucidrains/vit-pytorch ... check it out
for some einops/einsum fun
* Simple transformer style inspired by Andrej Karpathy's https://github.com/karpathy/minGPT
* Bert reference code checks against Huggingface Transformers and Tensorflow Bert

DeiT model defs and weights from https://github.com/facebookresearch/deit,
paper `DeiT: Data-efficient Image Transformers` - https://arxiv.org/abs/2012.12877

Hacked together by / Copyright 2020 Ross Wightman
"""

import math
import logging
from functools import partial
from collections import OrderedDict
import warnings

import torch
import torch.nn as nn
import torch.nn.functional as F

from timm.models.layers import DropPath, to_2tuple, trunc_normal_

from ..builder import ROTATED_BACKBONES

from mmengine.model import BaseModule
from mmengine import MessageHub

from mmengine.visualization.utils import (convert_overlay_heatmap, img_from_canvas)
from typing import Optional, Tuple
import cv2
import numpy as np

from .resnet import ResNet

from tensorboardX import SummaryWriter


_logger = logging.getLogger(__name__)

def draw_featmap(featmap: torch.Tensor,
                    overlaid_image: Optional[np.ndarray] = None,
                    channel_reduction: Optional[str] = 'squeeze_mean',
                    topk: int = 20,
                    arrangement: Tuple[int, int] = (4, 5),
                    resize_shape: Optional[tuple] = None,
                    alpha: float = 0.5) -> np.ndarray:
    """Draw featmap.

    - If `overlaid_image` is not None, the final output image will be the
        weighted sum of img and featmap.

    - If `resize_shape` is specified, `featmap` and `overlaid_image`
        are interpolated.

    - If `resize_shape` is None and `overlaid_image` is not None,
        the feature map will be interpolated to the spatial size of the image
        in the case where the spatial dimensions of `overlaid_image` and
        `featmap` are different.

    - If `channel_reduction` is "squeeze_mean" and "select_max",
        it will compress featmap to single channel image and weighted
        sum to `overlaid_image`.

    - If `channel_reduction` is None

        - If topk <= 0, featmap is assert to be one or three
        channel and treated as image and will be weighted sum
        to ``overlaid_image``.
        - If topk > 0, it will select topk channel to show by the sum of
        each channel. At the same time, you can specify the `arrangement`
        to set the window layout.

    Args:
        featmap (torch.Tensor): The featmap to draw which format is
            (C, H, W).
        overlaid_image (np.ndarray, optional): The overlaid image.
            Defaults to None.
        channel_reduction (str, optional): Reduce multiple channels to a
            single channel. The optional value is 'squeeze_mean'
            or 'select_max'. Defaults to 'squeeze_mean'.
        topk (int): If channel_reduction is not None and topk > 0,
            it will select topk channel to show by the sum of each channel.
            if topk <= 0, tensor_chw is assert to be one or three.
            Defaults to 20.
        arrangement (Tuple[int, int]): The arrangement of featmap when
            channel_reduction is None and topk > 0. Defaults to (4, 5).
        resize_shape (tuple, optional): The shape to scale the feature map.
            Defaults to None.
        alpha (Union[int, List[int]]): The transparency of featmap.
            Defaults to 0.5.

    Returns:
        np.ndarray: RGB image.
    """
    import matplotlib.pyplot as plt
    assert isinstance(featmap,
                        torch.Tensor), (f'`featmap` should be torch.Tensor,'
                                        f' but got {type(featmap)}')
    assert featmap.ndim == 3, f'Input dimension must be 3, ' \
                                f'but got {featmap.ndim}'
    featmap = featmap.detach().cpu()

    if overlaid_image is not None:
        if overlaid_image.ndim == 2:
            overlaid_image = cv2.cvtColor(overlaid_image,
                                            cv2.COLOR_GRAY2RGB)

        if overlaid_image.shape[:2] != featmap.shape[1:]:
            warnings.warn(
                f'Since the spatial dimensions of '
                f'overlaid_image: {overlaid_image.shape[:2]} and '
                f'featmap: {featmap.shape[1:]} are not same, '
                f'the feature map will be interpolated. '
                f'This may cause mismatch problems !')
            if resize_shape is None:
                featmap = F.interpolate(
                    featmap[None],
                    overlaid_image.shape[:2],
                    mode='bilinear',
                    align_corners=False)[0]

    if resize_shape is not None:
        featmap = F.interpolate(
            featmap[None],
            resize_shape,
            mode='bilinear',
            align_corners=False)[0]
        if overlaid_image is not None:
            overlaid_image = cv2.resize(overlaid_image, resize_shape[::-1])

    if channel_reduction is not None:
        assert channel_reduction in [
            'squeeze_mean', 'select_max'], \
            f'Mode only support "squeeze_mean", "select_max", ' \
            f'but got {channel_reduction}'
        if channel_reduction == 'select_max':
            sum_channel_featmap = torch.sum(featmap, dim=(1, 2))
            _, indices = torch.topk(sum_channel_featmap, 1)
            feat_map = featmap[indices]
        else:
            feat_map = torch.mean(featmap, dim=0)
        return convert_overlay_heatmap(feat_map, overlaid_image, alpha)
    elif topk <= 0:
        featmap_channel = featmap.shape[0]
        assert featmap_channel in [
            1, 3
        ], ('The input tensor channel dimension must be 1 or 3 '
            'when topk is less than 1, but the channel '
            f'dimension you input is {featmap_channel}, you can use the'
            ' channel_reduction parameter or set topk greater than '
            '0 to solve the error')
        return convert_overlay_heatmap(featmap, overlaid_image, alpha)
    else:
        row, col = arrangement
        channel, height, width = featmap.shape
        assert row * col >= topk, 'The product of row and col in ' \
                                    'the `arrangement` is less than ' \
                                    'topk, please set the ' \
                                    '`arrangement` correctly'

        # Extract the feature map of topk
        topk = min(channel, topk)
        sum_channel_featmap = torch.sum(featmap, dim=(1, 2))
        _, indices = torch.topk(sum_channel_featmap, topk)
        topk_featmap = featmap[indices]

        fig = plt.figure(frameon=False)
        # Set the window layout
        fig.subplots_adjust(
            left=0, right=1, bottom=0, top=1, wspace=0, hspace=0)
        dpi = fig.get_dpi()
        fig.set_size_inches((width * col + 1e-2) / dpi,
                            (height * row + 1e-2) / dpi)
        for i in range(topk):
            axes = fig.add_subplot(row, col, i + 1)
            axes.axis('off')
            axes.text(2, 15, f'channel: {indices[i]}', fontsize=10)
            axes.imshow(
                convert_overlay_heatmap(topk_featmap[i], overlaid_image,
                                        alpha))
        image = img_from_canvas(fig.canvas)
        plt.close(fig)
        return image

def batch_index_select(x, idx):
    if len(x.size()) == 3:
        B, N, C = x.size()
        N_new = idx.size(1)
        offset = torch.arange(B, dtype=torch.long, device=x.device).view(B, 1) * N
        idx = idx + offset
        out = x.reshape(B * N, C)[idx.reshape(-1)].reshape(B, N_new, C)
        return out
    elif len(x.size()) == 2:
        B, N = x.size()
        N_new = idx.size(1)
        offset = torch.arange(B, dtype=torch.long, device=x.device).view(B, 1) * N
        idx = idx + offset
        out = x.reshape(B * N)[idx.reshape(-1)].reshape(B, N_new)
        return out
    else:
        raise NotImplementedError
    
def restore_tokens(tokens_filtered, selected_index, original_length=1024):
    """
    将筛选后的tokens还原为原始形状，并用全零向量填充被丢弃的位置。

    参数:
    - tokens_filtered: 筛选后的tokens，形状为 [B, N, C]
    - selected_index: 选择的索引，形状为 [B, N, 1]
    - original_length: 原始tokens的长度，默认为1024

    返回:
    - restored_tokens: 还原后的tokens，形状为 [B, original_length, C]
    """
    B, N, C = tokens_filtered.shape

    # 初始化全零张量
    restored_tokens = torch.zeros((B, original_length, C), device=tokens_filtered.device, dtype=tokens_filtered.dtype)

    # 处理selected_index的形状，从 [B, N, 1] 转为 [B, N]
    selected_indices = selected_index.squeeze(-1).to(torch.long)  # [B, N]

    # 生成批次索引 [B, 1]，然后扩展为 [B, N] 以匹配selected_indices
    batch_indices = torch.arange(B, device=tokens_filtered.device).unsqueeze(1).expand(-1, N)  # [B, N]

    # 使用高级索引将筛选后的tokens放回到原始位置
    restored_tokens[batch_indices, selected_indices] = tokens_filtered

    return restored_tokens
    
class LayerNorm2d(nn.Module):
    def __init__(self, num_channels, eps=1e-5, elementwise_affine=True):
        super(LayerNorm2d, self).__init__()
        self.layer_norm = nn.LayerNorm(num_channels, eps, elementwise_affine)
    
    def forward(self, x):
        # x: [N, C, H, W]
        N, C, H, W = x.size()
        # 转换为 [N, H, W, C] 以应用 LayerNorm
        x = x.permute(0, 2, 3, 1)
        x = self.layer_norm(x)
        # 转换回 [N, C, H, W]
        x = x.permute(0, 3, 1, 2)
        return x

class Mlp(nn.Module):
    def __init__(
        self,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=nn.GELU,
        drop=0.0,
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class Attention(nn.Module):
    def __init__(
        self,
        dim,
        num_heads=8,
        qkv_bias=False,
        qk_scale=None,
        attn_drop=0.0,
        proj_drop=0.0,
    ):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
        self.scale = qk_scale or head_dim**-0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def softmax_with_policy(self, attn, policy, eps=1e-6):
        B, N, _ = policy.size()
        B, H, N, N = attn.size()
        attn_policy = policy.reshape(B, 1, 1, N)  # * policy.reshape(B, 1, N, 1)
        eye = torch.eye(N, dtype=attn_policy.dtype, device=attn_policy.device).view(
            1, 1, N, N
        )
        attn_policy = attn_policy + (1.0 - attn_policy) * eye
        max_att = torch.max(attn, dim=-1, keepdim=True)[0]
        attn = attn - max_att
        # attn = attn.exp_() * attn_policy
        # return attn / attn.sum(dim=-1, keepdim=True)

        # for stable training
        attn = attn.to(torch.float32).exp_() * attn_policy.to(torch.float32)
        attn = (attn + eps / N) / (attn.sum(dim=-1, keepdim=True) + eps)
        return attn.type_as(max_att)

    def forward(self, x, policy):
        B, N, C = x.shape
        qkv = (
            self.qkv(x)
            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
            .permute(2, 0, 3, 1, 4)
        )
        q, k, v = (
            qkv[0],
            qkv[1],
            qkv[2],
        )  # make torchscript happy (cannot use tensor as tuple)

        attn = (q @ k.transpose(-2, -1)) * self.scale

        if policy is None:
            attn = attn.softmax(dim=-1)
        else:
            attn = self.softmax_with_policy(attn, policy)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class Block(nn.Module):

    def __init__(
        self,
        dim,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=False,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            qk_scale=qk_scale,
            attn_drop=attn_drop,
            proj_drop=drop,
        )
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )

    def forward(self, x, policy=None):
        x = x + self.drop_path(self.attn(self.norm1(x), policy=policy))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x


class PatchEmbed(nn.Module):
    """Image to Patch Embedding"""

    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches

        self.proj = nn.Conv2d(
            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
        )

    def forward(self, x):
        B, C, H, W = x.shape
        # FIXME look at relaxing size constraints
        assert (
            H == self.img_size[0] and W == self.img_size[1]
        ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x


class HybridEmbed(nn.Module):
    """CNN Feature Map Embedding
    Extract feature map from CNN, flatten, project to embedding dim.
    """

    def __init__(
        self, backbone, img_size=224, feature_size=None, in_chans=3, embed_dim=768
    ):
        super().__init__()
        assert isinstance(backbone, nn.Module)
        img_size = to_2tuple(img_size)
        self.img_size = img_size
        self.backbone = backbone
        if feature_size is None:
            with torch.no_grad():
                # FIXME this is hacky, but most reliable way of determining the exact dim of the output feature
                # map for all networks, the feature metadata has reliable channel and stride info, but using
                # stride to calc feature dim requires info about padding of each stage that isn't captured.
                training = backbone.training
                if training:
                    backbone.eval()
                o = self.backbone(torch.zeros(1, in_chans, img_size[0], img_size[1]))
                if isinstance(o, (list, tuple)):
                    o = o[-1]  # last feature if backbone outputs list/tuple of features
                feature_size = o.shape[-2:]
                feature_dim = o.shape[1]
                backbone.train(training)
        else:
            feature_size = to_2tuple(feature_size)
            if hasattr(self.backbone, "feature_info"):
                feature_dim = self.backbone.feature_info.channels()[-1]
            else:
                feature_dim = self.backbone.num_features
        self.num_patches = feature_size[0] * feature_size[1]
        self.proj = nn.Conv2d(feature_dim, embed_dim, 1)

    def forward(self, x):
        x = self.backbone(x)
        if isinstance(x, (list, tuple)):
            x = x[-1]  # last feature if backbone outputs list/tuple of features
        x = self.proj(x).flatten(2).transpose(1, 2)
        return x
    
class DensityAwareModule(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
        super(DensityAwareModule, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding)
        self.bn = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        
    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x

class SparsificationPredictionModule(nn.Module):
    def __init__(self, embed_dim=384):
        super().__init__()
        self.in_conv = nn.Sequential(
            nn.LayerNorm(embed_dim), nn.Linear(embed_dim, embed_dim), nn.GELU()
        )

        self.out_conv = nn.Sequential(
            nn.LayerNorm(embed_dim),
            nn.Linear(embed_dim, embed_dim // 2),
            nn.GELU(),
            nn.Linear(embed_dim // 2, embed_dim // 4),
            nn.GELU(),
            nn.Linear(embed_dim // 4, 2),
            nn.LogSoftmax(dim=-1),
        )
        
        # self.dam = DensityAwareModule(embed_dim, embed_dim, kernel_size=3, stride=1, padding=1)

    def mask_generator(self, B, H, W, gt_bboxes):
        import cv2
        import numpy as np
        masks = []
        for i in range(B):
            # 初始化当前图像的mask
            mask = np.zeros((H, W), dtype=np.uint8)
            
            # 全部初始化为0.5
            mask.fill(0.8)
            
            # 取出第i张图像对应的真值框列表
            bboxes = gt_bboxes.data[0][i]
            
            if bboxes.ndim == 3:
                bboxes = bboxes.squeeze(0)
            
            for bbox in bboxes:
                # bbox: torch.tensor([cx, cy, w, h, angle])
                # 使用item()或float()从tensor中取出数值
                cx = float(bbox[0])
                cy = float(bbox[1])
                w = float(bbox[2])
                h = float(bbox[3])
                angle = float(bbox[4])
                
                # 解释变量作用：
                # cx, cy：旋转框中心点坐标(列、行)
                # w, h：旋转框宽高
                # angle：旋转角度(弧度)
                
                # 未旋转矩形的半宽半高
                half_w = w / 2.0
                half_h = h / 2.0
                
                # 矩形未旋转时的顶点 (相对中心点)
                # 顺序：P1, P2, P3, P4
                pts = np.array([
                    [-half_w, -half_h],
                    [-half_w,  half_h],
                    [ half_w,  half_h],
                    [ half_w, -half_h]
                ], dtype=np.float32)
                
                # 旋转矩形，使用旋转公式
                cos_theta = math.cos(angle)
                sin_theta = math.sin(angle)
                
                x_coords = pts[:, 0]
                y_coords = pts[:, 1]

                # 旋转后的点坐标
                x_rot = x_coords * cos_theta - y_coords * sin_theta
                y_rot = x_coords * sin_theta + y_coords * cos_theta
                
                # 平移到 (cx, cy)
                x_final = x_rot + cx
                y_final = y_rot + cy
                
                # 转换为int32的坐标，并匹配cv2.fillPoly函数要求的输入格式
                poly_pts = np.stack([x_final, y_final], axis=1).astype(np.int32)

                # 使用cv2.fillPoly填充
                cv2.fillPoly(mask, [poly_pts], 1)

            # 将生成的mask存入masks列表
            masks.append(mask)
        
        for i in range(len(masks)):
            masks[i] = torch.tensor(masks[i], dtype=torch.float32).unsqueeze(0)
            
        # 将masks列表转换为tensor
        masks = torch.stack(masks, dim=0)
        
        return masks


    def forward(self, x, policy):
        # x B, N, C
        B, N, C = x.size()
        message_hub = MessageHub.get_current_instance()
        gt_bboxes = message_hub.get_info('gt_bboxes')
        if self.training:
            masks = self.mask_generator(B, 32, 32, gt_bboxes).to(x.device) # B, 1, 32, 32
            masks = masks.squeeze(1) # B, 32, 32
            masks = masks.view(B, -1).unsqueeze(-1) # B, 1024， 1
            # B, N, C
            masks = masks.expand(B, N, C)
        else:
            masks = policy

        x = self.in_conv(x) # B, N, C
        local_x = x[:, :, : C // 2] # B, N, C//2
        global_x = (x[:, :, C // 2 : ] * policy).sum(dim=1, keepdim=True) / torch.sum(
            policy, dim=1, keepdim=True
        ) # B, 1, C//2
        x = torch.cat([local_x, global_x.expand(B, N, C // 2)], dim=-1) # B, N, C
        # if self.training:
        #     x = x * masks
        return self.out_conv(x) # B, N, 2
    
class BasicBlock(nn.Module):
    expansion = 1
    
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, 
                               padding=1, bias=False)
        self.ln1 = LayerNorm2d(out_channels) 
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, 
                               padding=1, bias=False)
        self.ln2 = LayerNorm2d(out_channels)  
        self.downsample = downsample
        self.relu = nn.ReLU(inplace=True)

        # 初始化权重
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, LayerNorm2d):
                nn.init.constant_(m.layer_norm.weight, 1)
                nn.init.constant_(m.layer_norm.bias, 0)

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.ln1(out)  # 使用 LayerNorm2d
        out = self.relu(out)
        
        out = self.conv2(out)
        out = self.ln2(out)  # 使用 LayerNorm2d
        
        if self.downsample is not None:
            identity = self.downsample(x)
        
        out += identity
        out = self.relu(out)
        return out
    
    
class MultiScaleFeatureBranch(nn.Module):
    def __init__(self, block, layers, channels = [16, 32, 64, 128, 256]):
        super(MultiScaleFeatureBranch, self).__init__()
        self.in_channels = channels[0]  
        self.conv1 = nn.Conv2d(3, channels[0], kernel_size=7, stride=1, padding=3, bias=False)  # 输出通道: 64, 尺寸: 512x512
        self.dropout = nn.Dropout(p=0.5)  # 添加Dropout层
        
        self.ln1 = LayerNorm2d(channels[0])  # 只对通道维度进行归一化
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)  # 输出尺寸: 256x256
        
        self.layer1 = self._make_layer(block, channels[0], layers[0], stride=1)    # 输出通道: 64, 尺寸: 256x256
        self.layer2 = self._make_layer(block, channels[1], layers[1], stride=2)   # 输出通道: 128, 尺寸: 128x128
        self.layer3 = self._make_layer(block, channels[2], layers[2], stride=2)   # 输出通道: 256, 尺寸: 64x64
        self.layer4 = self._make_layer(block, channels[3], layers[3], stride=2)   # 输出通道: 768, 尺寸: 32x32
        self.layer5 = self._make_layer(block, channels[4], layers[4], stride=2)   # 输出通道: 768, 尺寸: 16x16
        
        # 初始化权重
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, LayerNorm2d):
                nn.init.constant_(m.layer_norm.weight, 1)
                nn.init.constant_(m.layer_norm.bias, 0)

    def _make_layer(self, block, out_channels, blocks, stride=1):
        downsample = None
        if stride != 1 or self.in_channels != out_channels * block.expansion:
            downsample = nn.Sequential(
                nn.Conv2d(self.in_channels, out_channels * block.expansion, 
                          kernel_size=1, stride=stride, bias=False),
                LayerNorm2d(out_channels * block.expansion)  # 使用自定义的 LayerNorm2d
            )
        
        layers = []
        layers.append(block(self.in_channels, out_channels, stride, downsample))
        self.in_channels = out_channels * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.in_channels, out_channels))
        
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.conv1(x)
        x = self.ln1(x)  # 已在 LayerNorm2d 中处理维度转换
        x = F.relu(x, inplace=True)
        x = self.dropout(x)  # 在激活函数之后添加Dropout
        
        x = self.maxpool(x)
        outs = []
        
        x = self.layer1(x)
        outs.append(x)
        x = self.layer2(x)
        outs.append(x)
        x = self.layer3(x)
        outs.append(x)
        x = self.layer4(x)
        outs.append(x)
        x = self.layer5(x)
        outs.append(x)

        return outs
    

class FusionModule(nn.Module):
    def __init__(self, conv_channels=[16, 32, 64, 128, 256], attention_channels=128, normalize_after_add=True):
        super(FusionModule, self).__init__()
        self.attention_channels = attention_channels
        self.normalize_after_add = normalize_after_add

        # Filter out channels that need adjustment (not equal to attention_channels)
        self.adjusted_channels = [c + attention_channels for c in conv_channels]

        # Initialize LayerNorm2d for each adjusted channel
        # self.attention_norms = nn.ModuleList([
        #     LayerNorm2d(c) for c in self.adjusted_channels
        # ])

        if self.normalize_after_add:
            self.add_norms = nn.ModuleList([
                LayerNorm2d(c) for c in self.adjusted_channels
            ])

        # Initialize weights for LayerNorm2d
        for m in self.modules():
            if isinstance(m, LayerNorm2d):
                nn.init.constant_(m.layer_norm.weight, 1)
                nn.init.constant_(m.layer_norm.bias, 0)

    def forward(self, conv_features, attention_feature):
        """
        Args:
            conv_features (List[Tensor]): List of convolutional feature maps.
            attention_feature (Tensor): Attention feature map.

        Returns:
            List[Tensor]: List of fused feature maps.
        """
        added_features = []
        norm_idx = 0  # Index for attention_norms and add_norms

        for i, conv in enumerate(conv_features):
            # Get current convolutional feature dimensions
            _, C, H, W = conv.size()

            # # Adjust attention feature channels if necessary
            # if C != self.attention_channels:
            #     if self.attention_channels % C != 0:
            #         raise ValueError(f"attention_channels ({self.attention_channels}) is not divisible by conv_channels ({C})")

            #     group = self.attention_channels // C
            #     # Reshape and average to reduce channels
            #     att = attention_feature.view(N, C, group, H_att, W_att).mean(dim=2)  # (N, C, H_att, W_att)
                
            #     # Apply LayerNorm2d
            #     att = self.attention_norms[norm_idx](att)
            # else:
            #     att = attention_feature  # No adjustment needed
            att = attention_feature

            # Upsample attention feature to match spatial dimensions of conv feature
            att = F.interpolate(att, size=(H, W), mode='bilinear', align_corners=False)

            # Add attention feature to convolutional feature
            # added = conv + att
            added = torch.cat([conv, att], dim=1)

            # Optionally apply LayerNorm2d after addition
            # if self.normalize_after_add and C != self.attention_channels:
            added = self.add_norms[norm_idx](added)

            added_features.append(added)

            # if C != self.attention_channels:
            #     norm_idx += 1  # Increment index only when adjustment was applied
            norm_idx += 1

        return added_features
    

@ROTATED_BACKBONES.register_module()
class DenSeAdViT(BaseModule):
    """Vision Transformer
    DensitySensitiveVisionTransformerwithAdaptiveTokens

    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`  -
        https://arxiv.org/abs/2010.11929
    """

    def __init__(
        self,
        img_size=224,
        patch_size=16,
        in_chans=3,
        embed_dim=768,
        depth=12,
        num_heads=12,
        mlp_ratio=4.0,
        qkv_bias=True,
        qk_scale=None,  # zhaobudao
        representation_size=None,  # wu
        drop_rate=0.0,  # 0.0
        attn_drop_rate=0.0,  #  wu
        drop_path_rate=0.0,  #  .2 vit-b 0.5
        hybrid_backbone=None,  # wu
        norm_layer=None,
        pruning_loc=None,  # gezhongdouyou [4,8,12]
        token_ratio=None,  # [0.7, 0.7**2, 0.7**3]
        distill=False,  # True
    ):
        """
        Args:
            img_size (int, tuple): input image size
            patch_size (int, tuple): patch size
            in_chans (int): number of input channels
            num_classes (int): number of classes for classification head
            embed_dim (int): embedding dimension
            depth (int): depth of transformer
            num_heads (int): number of attention heads
            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
            qkv_bias (bool): enable bias for qkv if True
            qk_scale (float): override default qk scale of head_dim ** -0.5 if set
            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
            drop_rate (float): dropout rate
            attn_drop_rate (float): attention dropout rate
            drop_path_rate (float): stochastic depth rate
            hybrid_backbone (nn.Module): CNN backbone to use in-place of PatchEmbed module
            norm_layer: (nn.Module): normalization layer
        """
        super().__init__()

        print("## diff vit pruning method")
        self.num_features = self.embed_dim = (
            embed_dim  # num_features for consistency with other models
        )
        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)

        if hybrid_backbone is not None:
            self.patch_embed = HybridEmbed(
                hybrid_backbone,
                img_size=img_size,
                in_chans=in_chans,
                embed_dim=embed_dim,
            )
        else:
            self.patch_embed = PatchEmbed(
                img_size=img_size,
                patch_size=patch_size,
                in_chans=in_chans,
                embed_dim=embed_dim,
            )
        num_patches = self.patch_embed.num_patches

        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
        self.pos_drop = nn.Dropout(p=drop_rate)
        
        # self.multi_scale_feature_branch = MultiScaleFeatureBranch(BasicBlock, [3, 6, 4, 3, 1], channels=conv_channels)
        self.multi_scale_feature_branch = ResNet(depth=34)

        dpr = [
            x.item() for x in torch.linspace(0, drop_path_rate, depth)
        ]  # stochastic depth decay rule
        self.blocks = nn.ModuleList(
            [
                Block(
                    dim=embed_dim,
                    num_heads=num_heads,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    qk_scale=qk_scale,
                    drop=drop_rate,
                    attn_drop=attn_drop_rate,
                    drop_path=dpr[i],
                    norm_layer=norm_layer,
                )
                for i in range(depth)
            ]
        )
        self.norm = norm_layer(embed_dim)

        # Representation layer
        if representation_size:
            self.num_features = representation_size
            self.pre_logits = nn.Sequential(
                OrderedDict(
                    [
                        ("fc", nn.Linear(embed_dim, representation_size)),
                        ("act", nn.Tanh()),
                    ]
                )
            )
        else:
            self.pre_logits = nn.Identity()

        predictor_list = [SparsificationPredictionModule(embed_dim) for _ in range(len(pruning_loc))]

        self.score_predictor = nn.ModuleList(predictor_list)

        self.distill = distill

        self.pruning_loc = pruning_loc
        self.token_ratio = token_ratio
        
        self.conv_channels = [64, 128, 256, 512]
        
        self.fusion_module = FusionModule(conv_channels=self.conv_channels, attention_channels=embed_dim)
        
        # Initialize TensorBoard SummaryWriter
        self.writer = SummaryWriter(log_dir='./logs')
        self.global_step = 0  # Initialize a step counter

        trunc_normal_(self.pos_embed, std=0.02)
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    @torch.jit.ignore
    def no_weight_decay(self):
        return {"pos_embed"}
    
    def forward(self, x):
        B = x.shape[0]  # torch.Size([2, 3, 512, 512])
        ms_fs = self.multi_scale_feature_branch(x)
        # torch.Size([2, 64, 128, 128])
        # torch.Size([2, 128, 64, 64])
        # torch.Size([2, 256, 32, 32]) 
        # torch.Size([2, 512, 16, 16])
        x = self.patch_embed(
            x
        )  # torch.Size([2, 1024, 768]) 512/16 = 32 32*32 = 1024 embedding dim = 768

        x = x + self.pos_embed  # torch.Size([2, 1024, 768])
        x = self.pos_drop(x)  # torch.Size([2, 1024, 768])

        p_count = 0
        out_pred_prob = []
        init_n = 32 * 32  # init_n = 32 * 32
        prev_decision = torch.ones(
            B, init_n, 1, dtype=x.dtype, device=x.device
        )  # torch.Size([2, 1024, 1])
        policy = torch.ones(
            B, init_n, 1, dtype=x.dtype, device=x.device
        )  # torch.Size([2, 1024, 1])
        # selected_index = torch.arange(0, 1024, dtype=torch.int32, device=x.device).unsqueeze(0).repeat(B, 1)

        for i, blk in enumerate(self.blocks):
            if i in self.pruning_loc:
                spatial_x = x  # 直接使用 x
                pred_score = self.score_predictor[p_count](
                    spatial_x, prev_decision
                ).reshape(
                    B, -1, 2
                )  # torch.Size([2, 1024, 2])

                # if self.training:
                hard_keep_decision = (
                    F.gumbel_softmax(pred_score, hard=True)[:, :, 0:1]
                    * prev_decision
                )  # torch.Size([2, 1024, 1])
                out_pred_prob.append(
                    hard_keep_decision.reshape(B, init_n)
                )  # [torch.Size([2, 1024]), ]
                policy = hard_keep_decision  # torch.Size([2, 1024, 1])
                # x = blk(x, policy=policy)
                x = blk(x)
                prev_decision = hard_keep_decision  # torch.Size([2, 1024, 1])

                # else:
                #     score = pred_score[:, :, 0] # torch.Size([B, 1024])
                #     num_keep_node = int(init_n * self.token_ratio[p_count]) # 32*32*0.7 716
                #     keep_policy = torch.argsort(score, dim=1, descending=True)[
                #         :, :num_keep_node
                #     ]
                #     now_policy = keep_policy  # 直接使用 keep_policy B, num_keep_node
                #     x = batch_index_select(x, now_policy) # torch.Size([2, 716, 768]) B  num_keep_node  768
                #     selected_index = batch_index_select(selected_index, now_policy)
                #     prev_decision = batch_index_select(prev_decision, keep_policy)
                #     x = blk(x)
                p_count += 1
            else:
                if self.training:
                    # x = blk(x, policy)
                    x = blk(x)
                else:
                    x = blk(x)

        x = self.norm(x)  # torch.Size([2, 1024, 768])
        features = x
        # if  self.training:
        #     features = x  # torch.Size([2, 1024, 768])
        # else:
        #     features = restore_tokens(x, selected_index, init_n)
        features = features.permute(0, 2, 1)  # [2, 768, 1024]
        # [2, 768, 1024] -> [2, 768, 32, 32]
        features = features.reshape(B, features.size(1), 32, 32)
        
        features = self.fusion_module(ms_fs, features)
        
        # Log feature maps to TensorBoard
        for idx, feature in enumerate(features):
            fmap = draw_featmap(feature[0])  # Convert to numpy array (H, W, 3)
            
            # Normalize the image to [0, 1] if it's not already
            fmap = fmap.astype(np.float32) / 255.0 if fmap.max() > 1.0 else fmap.astype(np.float32)
            
            # Convert numpy image (H, W, C) to torch tensor (C, H, W)
            fmap_tensor = torch.from_numpy(fmap).permute(2, 0, 1)
            
            # Add the image to TensorBoard
            self.writer.add_image(f'FeatureMap_{idx}', fmap_tensor, self.global_step)
        
        # Increment the global step counter
        self.global_step += 1
        
        return features
